﻿/*  Copyright (C) 2012 AmaroK86 (marcidm 'at' hotmail 'dot' com)
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.

 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.

 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
using System;
using System.IO;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using ICSharpCode.SharpZipLib.Zip.Compression.Streams;
using Gibbed.IO;

namespace AmaroK86.MassEffect3.ZlibBlock
{
    public static class ZBlock
    {
        public static readonly uint magic = 0x9E2A83C1;
        public static readonly uint maxSegmentSize = 0x20000;

        /// <summary>
        ///     compress a part of the byte array into a Zlib Block.
        /// </summary>
        /// <param name="buffer">byte array to compress</param>
        /// <param name="offset">starting offset inside the array</param>
        /// <param name="count">num of bytes to compress starting from the offset</param>
        /// <returns>a compressed byte array block.</returns>
        public static byte[] Compress(byte[] buffer, int offset, int count)
        {
            if(buffer == null)
                throw new ArgumentNullException();
            if (count < 0)
                throw new FormatException();
            if (offset + count > buffer.Length)
                throw new IndexOutOfRangeException();

            byte[] finalBlock;
            using (MemoryStream headBlock = new MemoryStream(), dataBlock = new MemoryStream())
            {
                DeflaterOutputStream zipStream;

                int numSeg = (int)Math.Ceiling((double)count / (double)maxSegmentSize);

                headBlock.WriteValueU32(magic);
                headBlock.WriteValueU32(maxSegmentSize);
                headBlock.WriteValueU32(0x0);            //total compressed size, still to calculate
                headBlock.WriteValueS32(count);          //total uncompressed size

                for (int i = count; i > 0; i -= (int)maxSegmentSize)
                {
                    int copyBytes = Math.Min(i, (int)maxSegmentSize);
                    uint precCompSize = (uint)dataBlock.Length;
                    zipStream = new DeflaterOutputStream(dataBlock);
                    zipStream.Write(buffer, offset + (count - i), copyBytes);
                    zipStream.Flush();
                    zipStream.Finish();
                    headBlock.WriteValueU32((uint)dataBlock.Length - precCompSize); //compressed segment size
                    headBlock.WriteValueS32(copyBytes); //uncompressed segment size
                    //Console.WriteLine("  Segment size: {0}, total read: {1}, compr size: {2}", maxSegmentSize, copyBytes, (uint)dataBlock.Length - precCompSize);
                }

                headBlock.Seek(8, SeekOrigin.Begin);
                headBlock.WriteValueS32((int)dataBlock.Length); // total compressed size

                finalBlock = new byte[headBlock.Length + dataBlock.Length];
                Buffer.BlockCopy(headBlock.ToArray(), 0, finalBlock, 0, (int)headBlock.Length);
                Buffer.BlockCopy(dataBlock.ToArray(), 0, finalBlock, (int)headBlock.Length, (int)dataBlock.Length);
            }
            return finalBlock;
        }

        /// <summary>
        ///     compress an entire byte array into a Zlib Block.
        /// </summary>
        /// <param name="buffer">byte array to compress</param>
        /// <returns>a compressed byte array block.</returns>
        public static byte[] Compress(byte[] buffer)
        {
            if (buffer == null)
                throw new ArgumentNullException();
            return Compress(buffer, 0, buffer.Length);
        }

        /// <summary>
        ///     compress a part of the stream into a Zlib Block starting from the current stream position
        /// </summary>
        /// <param name="inStream">Stream to compress</param>
        /// <param name="count">num of bytes to compress starting from the Stream position</param>
        /// <returns>a compressed byte array block.</returns>
        public static byte[] Compress(Stream inStream, int count)
        {
            if (count < 0)
                throw new FormatException();
            if (inStream.Position + count > inStream.Length)
                throw new ArgumentOutOfRangeException();
            byte[] buffer = new byte[count];
            inStream.Read(buffer, 0, count);
            return Compress(buffer,0,count);
        }

        /// <summary>
        ///     decompress an entire byte array from a Zlib Block.
        /// </summary>
        /// <param name="buffer">byte array to decompress</param>
        /// <returns>a decompressed byte array.</returns>
        public static byte[] Decompress(byte[] buffer)
        {
            if (buffer == null)
                throw new ArgumentNullException();
            return Decompress(buffer, 0, buffer.Length);
        }

        /// <summary>
        ///     decompress a Zlib Block from a part of byte array.
        /// </summary>
        /// <param name="buffer">byte array to decompress</param>
        /// <param name="offset">starting offset inside the array</param>
        /// <param name="count">num of bytes to decompress starting from the offset</param>
        /// <returns>a decompressed byte array block.</returns>
        public static byte[] Decompress(byte[] buffer, int offset, int count)
        {
            if (buffer == null)
                throw new ArgumentNullException();
            if (count < 0)
                throw new FormatException();
            if (offset + count > buffer.Length)
                throw new IndexOutOfRangeException();

            byte[] outputBuffer;
            using (MemoryStream buffStream = new MemoryStream(buffer, offset, count))
            {
                InflaterInputStream zipStream;
                uint magicStream = buffStream.ReadValueU32();
                if (magicStream != magic && magicStream.Swap() != magic)
                {
                    throw new InvalidDataException("not a valid zlib block");
                }

                uint buffMaxSegmentSize = buffStream.ReadValueU32();
                if (buffMaxSegmentSize != maxSegmentSize)
                {
                    throw new FormatException();
                }

                uint totComprSize = buffStream.ReadValueU32();
                uint totUncomprSize = buffStream.ReadValueU32();

                outputBuffer = new byte[totUncomprSize];
                int numOfSegm = (int)Math.Ceiling((double)totUncomprSize / (double)maxSegmentSize);
                int headSegm = 16;
                int dataSegm = headSegm + (numOfSegm * 8);
                int buffOff = 0;

                for (int i = 0; i < numOfSegm; i++)
                {
                    buffStream.Seek(headSegm, SeekOrigin.Begin);
                    int comprSegm = buffStream.ReadValueS32();
                    int uncomprSegm = buffStream.ReadValueS32();
                    headSegm = (int)buffStream.Position;

                    buffStream.Seek(dataSegm, SeekOrigin.Begin);
                    //Console.WriteLine("compr size: {0}, uncompr size: {1}, data offset: 0x{2:X8}", comprSegm, uncomprSegm, dataSegm);
                    zipStream = new InflaterInputStream(buffStream);
                    zipStream.Read(outputBuffer, buffOff, uncomprSegm);
                    zipStream.Flush();
                    buffOff += uncomprSegm;
                    dataSegm += comprSegm;
                }
            }
            return outputBuffer;
        }

        /// <summary>
        ///     decompress a Zlib Block from a stream starting from the current stream position
        /// </summary>
        /// <param name="inStream">Stream to decompress</param>
        /// <param name="count">num of bytes to decompress starting from the Stream position</param>
        /// <returns>a decompressed byte array.</returns>
        public static byte[] Decompress(Stream inStream, int count)
        {
            if (count < 0)
                throw new FormatException();
            if (inStream.Position + count > inStream.Length)
                throw new ArgumentOutOfRangeException();
            byte[] buffer = new byte[count];
            inStream.Read(buffer, 0, count);
            return Decompress(buffer, 0, count);
        }
    }
}
